package cn.edu.recommender

import org.apache.spark.sql.SparkSession

import java.lang.Thread.sleep

object SameUserProductRecommender {

  // 定义常量
  val MONGODB_RATING_COLLECTION = "Rating"
  val MAX_RECOMMENDATION = 20

  val SAME_USER_PRODUCT_RECS = "SameUserProductRecs"


  def main(args: Array[String]): Unit = {

    // 定义配置
    val config = Map(
      "spark.cores" -> "local[*]",
      "mongo.uri" -> "mongodb://localhost:27017/recommender",
      "mongo.db" -> "recommender"
    )

    val mongoConfig = MongoConfig(config("mongo.uri"), config("mongo.db"))

    // 创建一个 SparkSession
    val spark = SparkSession.builder().appName("SameUserProductRecommender").master(config("spark.cores")).getOrCreate()

    // 在对 DataFrame 和 Dataset 进行操作许多操作都需要这个包进行支持
    import spark.implicits._

    //数据加载进来
    val ratingDF = spark
      .read
      .option("uri", mongoConfig.uri)
      .option("collection", MONGODB_RATING_COLLECTION)
      .format("com.mongodb.spark.sql")
      .load()
      .as[ProductRating]
      .map(rating => (rating.userId, rating.productId, rating.score))
      .toDF("userId", "productId", "rating")

    // 统计每个商品的评分个数，并通过内连接添加到 ratingDS 中
    val numRatersPerProduct = ratingDF.groupBy("productId").count()
    val ratingWithCountDF = ratingDF.join(numRatersPerProduct, "productId")

    // 将商品评分按 userId 两两配对，可以统计两个商品被同一用户做出评分的次数
    ratingWithCountDF.join(ratingWithCountDF, "userId")
      .toDF("userId", "product1", "rating1", "count1", "product2", "rating2", "count2")
      .select("userId", "product1", "count1", "product2", "count2")
      .createOrReplaceTempView("a")

    val cooccurrenceDF = spark.sql(
      """
        |select product1, product2, count(userId) as coocount,
        |first(count1) as count1, first(count2) as count2
        |from a
        |where product1 != product2
        |group by product1, product2
      """.stripMargin
    ).cache()

    // 用同现的次数和各自的次数，计算同现相似度
    val simDS = cooccurrenceDF.map {
      row =>
        val coocSim = row.getAs[Long]("coocount") / math.sqrt(row.getAs[Long]("count1") * row.getAs[Long]("count2"))
        (row.getAs[Int]("product1"), (row.getAs[Int]("product2"), coocSim))
    }
      .rdd
      .groupByKey()
      .map{
        case (productId, recs) =>
          ProductRecs(productId, recs.toList.sortWith(_._2>_._2).map(x=>Recommendation(x._1,x._2)).take(MAX_RECOMMENDATION))
      }
      .toDS()


    simDS
      .write
      .option("uri", mongoConfig.uri)
      .option("collection", SAME_USER_PRODUCT_RECS)
      .mode("overwrite")
      .format("com.mongodb.spark.sql")
      .save()


    sleep(300000)
    // 关闭 Spark
    spark.stop()
  }
}
